#include "../../DWMRICore/Fiber.h"
#include "../../DWMRICore/Scalar.h"
#include "../../DWMRICore/DTI.h"

#include "../CUDA_ComputeStructureTensor.h"
#include "CUDA_ST_Method_EuclideanVector.h"

#include <omp.h>


extern "C" void CUDA_FiberParameterization_EV_DTI(CFiber **fibers, 
	Fiber_EV_Param *parameters, const int w, const int h, const int length);
extern "C" void CUDA_FiberParameterization_EV_HARDI(CFiber **fibers, 
	int *count, Fiber_EV_Param *parameters, const int w, const int h, const int length);


void ComputeTensorEV(float *function, float *tensor)
{
	tensor[0] = 0.0f;
	tensor[1] = 0.0f;
	tensor[2] = 0.0f;
	tensor[3] = 0.0f;
	tensor[4] = 0.0f;
	tensor[5] = 0.0f;

	for (int i = 0; i < 9; ++i) {
		tensor[0] += function[i*3+0] * function[i*3+0];
		tensor[1] += function[i*3+0] * function[i*3+1];
		tensor[2] += function[i*3+0] * function[i*3+2];
		tensor[3] += function[i*3+1] * function[i*3+1];
		tensor[4] += function[i*3+1] * function[i*3+2];
		tensor[5] += function[i*3+2] * function[i*3+2];
	}
}

void StructureTensorEV(Fiber_EV_Param **fibers, float *tensor)
{
	double A[81];
	for (int k = 0; k < 27; ++k) {
		A[k*3+0] = fibers[k]->seed.pos.x;
		A[k*3+1] = fibers[k]->seed.pos.y;
		A[k*3+2] = fibers[k]->seed.pos.z;
	}

	/* compute pseudo inverse of matrix A */
	double AA[9];
	for (int i = 0; i < 3; ++i) {
		for (int j = 0; j < 3; ++j) {
			AA[i*3+j] = 0.0;
			for (int k = 0; k < 27; ++k) {
				AA[i*3+j] += A[k*3+j] * A[k*3+i];
			}
		}
	}
	double invAA[9];
	Inv33(AA, invAA);

	double pinvA[81];
	for (int i = 0; i < 3; ++i) {
		for (int j = 0; j < 27; ++j) {
			pinvA[i*27+j] = 0.0f;
			for (int k = 0; k < 3; ++k) {
				pinvA[i*27+j] += invAA[i*3+k] * A[j*3+k];
			}
		}
	}

	/* estimate functions */
	int flags[27];
	for (int i = 0; i < 27; ++i) {
		if ((fibers[i]->seed.dir.x * fibers[0]->seed.dir.x + fibers[i]->seed.dir.y * fibers[0]->seed.dir.y + fibers[i]->seed.dir.z * fibers[0]->seed.dir.z) > 0.0f) {
			flags[i] = 1;
		} else {
			flags[i] = 0;
		}
	}

	float *functions = new float[27];

	for (int i = 0; i < 9; ++i) {
		float y[27];

		if (i < 3) {
			for (int k = 0; k < 27; ++k) {
				if (flags[k] == 1) {
					y[k] = fibers[k]->parameters[i];
				} else {
					y[k] = -fibers[k]->parameters[i];
				}
			}
		} else {
			for (int k = 0; k < 27; ++k) {
				y[k] = fibers[k]->parameters[i];
			}
		}
		
		float x[3] = {0.0f, 0.0f, 0.0f};
		for (int k = 0; k < 27; ++k) {
			x[0] += (float)pinvA[k] * y[k];
			x[1] += (float)pinvA[27+k] * y[k];
			x[2] += (float)pinvA[54+k] * y[k];
		}
		functions[i*3+0] = x[0];
		functions[i*3+1] = x[1];
		functions[i*3+2] = x[2];
	}

	ComputeTensorEV(functions, tensor);

	delete[] functions;

}

void CPU_ComputeStructureTensor_EV_DTI(Fiber_EV_Param *prev, Fiber_EV_Param *curr, 
		Fiber_EV_Param *next, float *tensor, const int w, const int h, const int d)
{
	int offset[9];
	offset[ 0] = 0;
	offset[ 1] = -1;
	offset[ 2] = 1;
	offset[ 3] = -w-1;
	offset[ 4] = -w;
	offset[ 5] = -w+1;
	offset[ 6] = w-1;
	offset[ 7] = w;
	offset[ 8] = w+1;

	omp_set_num_threads(64);
	for (int y = 1; y < h - 1; ++y) {
#pragma omp parallel for            
		for (int x = 1; x < w - 1; ++x) {
			int index = y * w + x;
			Fiber_EV_Param *fibers[27];
			for (int i = 0; i < 9; ++i) {
				fibers[i] = &(curr[index+offset[i]]);
				fibers[i+9] = &(prev[index+offset[i]]);
				fibers[i+18] = &(next[index+offset[i]]);
			}

			float t[6] = {1.0f, 0.0f, 0.0f, 1.0f, 0.0f, 1.0f};
			StructureTensorEV(fibers, t);
			tensor[index*10+0]  = 1.0f;
			tensor[index*10+1]  = fibers[0]->seed.dir.x;
			tensor[index*10+2]  = fibers[0]->seed.dir.y;
			tensor[index*10+3]  = fibers[0]->seed.dir.z;
			tensor[index*10+4] = t[0];
			tensor[index*10+5] = t[1];
			tensor[index*10+6] = t[2];
			tensor[index*10+7] = t[3];
			tensor[index*10+8] = t[4];
			tensor[index*10+9] = t[5];
		}
	}
}


void CPU_ComputeStructureTensor_EV_HARDI(Fiber_EV_Param **paramters, int **counts, float *tensor,
	const int w, const int h, const int d)
{
	int offset[9];
	offset[ 0] = 0;
	offset[ 1] = -1;
	offset[ 2] = 1;
	offset[ 3] = -w-1;
	offset[ 4] = -w;
	offset[ 5] = -w+1;
	offset[ 6] = w-1;
	offset[ 7] = w;
	offset[ 8] = w+1;

	omp_set_num_threads(64);
	for (int y = 1; y < h - 1; ++y) {
#pragma omp parallel for            
		for (int x = 1; x < w - 1; ++x) {
			int index = y * w + x;
			for (int k = 0; k < counts[1][index]; ++k) {
				Fiber_EV_Param *fibers[27];

				fibers[0] = &(paramters[1][index*3+k]);
				float3 dir = fibers[0]->seed.dir;
				for (int l = 0; l < 3; ++l) {
					for (int i = 0; i < 9; ++i) {
						if (l == 1 && i == 0)
							continue;

						float max_v = -1.0f;
						int max_index;
						int temp_index = index + offset[i];
						int size = counts[l][temp_index];
						for (int j = 0; j < size; ++j) {
							float3 temp_dir = paramters[l][temp_index*3+j].seed.dir;
							float v = dir.x * temp_dir.x + dir.y * temp_dir.y + dir.z * temp_dir.z;
							if (fabs(v) > max_v) {
								max_v = fabs(v);
								max_index = j;
								if (l == 0) {
									fibers[9+i] = &(paramters[l][temp_index*3+j]);
								} else if (l == 1) {
									fibers[i] = &(paramters[l][temp_index*3+j]);
								} else {
									fibers[18+i] = &(paramters[l][temp_index*3+j]);
								}
							}               
						}
					}
				}

				float t[6] = {1.0f, 0.0f, 0.0f, 1.0f, 0.0f, 1.0f};
				StructureTensorEV(fibers, t);
				int id = (y * w + x) * 3 + k;
				tensor[id*10+0] = 1.0f;
				tensor[id*10+1] = dir.x;
				tensor[id*10+2] = dir.y;
				tensor[id*10+3] = dir.z;
				tensor[id*10+4] = t[0];
				tensor[id*10+5] = t[1];
				tensor[id*10+6] = t[2];
				tensor[id*10+7] = t[3];
				tensor[id*10+8] = t[4];
				tensor[id*10+9] = t[5];
			}
		}
	}
}

void ComputeStructureTensorMethodEV(const float fiber_length, const float maxLength, 
	const float stepSize, const int w, const int h, const int d, const int scale, const int mode, 
	char **confPathname, char **vertPathname, char *tempPathname)
{
	float length = fiber_length > maxLength ? maxLength : fiber_length;
	length = length / stepSize;

	float *tensor = NULL;
	if (mode == 0) {
		tensor = new float[w*h*d*10];
		for (int z = 0; z < d; ++z) {
			for (int y = 0; y < h; ++y) {
				for (int x = 0; x < w; ++x) {
					int index = (z * h + y) * w + x;
					tensor[index*10+0] = 1.0f;
					tensor[index*10+1] = 0.0f; tensor[index*10+2] = 0.0f; tensor[index*10+3] = 0.0f;
					tensor[index*10+4] = 1.0f; tensor[index*10+5] = 0.0f; tensor[index*10+6] = 0.0f;
					tensor[index*10+7] = 1.0f; tensor[index*10+8] = 0.0f; tensor[index*10+9] = 1.0f;
				}
			}
		}
	} else if (mode == 1) {
		tensor = new float[w*h*d*30];
		memset(tensor, 0, sizeof(float)*w*h*d*30);
		for (int z = 0; z < d; ++z) {
			for (int y = 0; y < h; ++y) {
				for (int x = 0; x < w; ++x) {
					int index = (z * h + y) * w + x;
					tensor[index*30+0] = 1.0f;
					tensor[index*30+1] = 0.0f; tensor[index*30+2] = 0.0f; tensor[index*30+3] = 0.0f;
					tensor[index*30+4] = 1.0f; tensor[index*30+5] = 0.0f; tensor[index*30+6] = 0.0f;
					tensor[index*30+7] = 1.0f; tensor[index*30+8] = 0.0f; tensor[index*30+9] = 1.0f;
				}
			}
		}
	} else {
		;
	}

	printf("\n=============== Begin Computing Structuren Tensor ===============\n");
	printf("=============== Method Euclidean Vector ===============\n");
	for (int z = 1; z < d - 1; ++z) {
		unsigned int hTimer;
		cutilCheckError(cutCreateTimer(&hTimer));
		cutilCheckError(cutResetTimer(hTimer));
		cutilCheckError(cutStartTimer(hTimer));
		printf("Slice z=%.04f ", (float)z/(float)scale);
		fflush(stdout);

		CFiber **fibers[3];
        int *count[3];
        int cw, ch;
		Fiber_EV_Param *parameters[3];

		for (int k = 0; k < 3; ++k) {
			ReadFibers(confPathname[z+k-1], vertPathname[z+k-1], &(fibers[k]), &(count[k]), cw, ch);

			if (mode == 0) {				/* dti */
				parameters[k] = new Fiber_EV_Param[cw*ch];
				memset(parameters[k], 0, sizeof(Fiber_EV_Param)*cw*ch);
				CUDA_FiberParameterization_EV_DTI(fibers[k], parameters[k], cw, ch, (int)length);
			} else if (mode == 1) {			/* hardi */
				parameters[k] = new Fiber_EV_Param[cw*ch*3];
				memset(parameters[k], 0, sizeof(Fiber_EV_Param)*cw*ch*3);
				CUDA_FiberParameterization_EV_HARDI(fibers[k], count[k], parameters[k], cw, ch, (int)length);
			} else {
				;
			}

			for (int i = 0; i < w * h; ++i) {
				SafeDeleteArray(fibers[k][i]);
			}
			SafeDeleteArray(fibers[k]);
			
		}
		if (mode == 0) {			/* dti */
			CPU_ComputeStructureTensor_EV_DTI(parameters[0], parameters[1], parameters[2], 
				&(tensor[z*w*h*10]), w, h, d);
		} else if (mode == 1) {		/* hardi */
			CPU_ComputeStructureTensor_EV_HARDI(parameters, count, &(tensor[z*w*h*30]), w, h, d);
		} else {
			;
		}

		for (int k = 0; k < 3; ++k) {
			SafeDeleteArray(parameters[k]);
			SafeDeleteArray(count[k]);
		}

		printf(" Done. Processing time: %f (sec)\n", cutGetTimerValue(hTimer)/1000);
	}
	printf("=============== Finish ===============\n\n");

	CDTI *pTempTensor = NULL;
	if (mode == 0) {
		pTempTensor = new CDTI();
		pTempTensor->CreateDTI(w, h, d);
		for (int z = 0; z < d; ++z) {
			for (int y = 0; y < h; ++y) {
				for (int x = 0; x < w; ++x) {
					int index = (z * h + y) * w + x;
					float dti[7] = {1.0f, 1.0f, 0.0f, 0.0f, 1.0f, 0.0f, 1.0f};
					dti[1] = tensor[index*10+4];
					dti[2] = tensor[index*10+5];
					dti[3] = tensor[index*10+6];
					dti[4] = tensor[index*10+7];
					dti[5] = tensor[index*10+8];
					dti[6] = tensor[index*10+9];
					pTempTensor->SetDTI(x, y, z, dti);
				}
			}
		}
		pTempTensor->SaveDTIFile(tempPathname);
		SafeDelete(pTempTensor);
	} else {
		pTempTensor = new CDTI();
		pTempTensor->CreateDTI(w, h, d);
		for (int z = 0; z < d; ++z) {
			for (int y = 0; y < h; ++y) {
				for (int x = 0; x < w; ++x) {
					int index = (z * h + y) * w + x;
					float dti[7] = {1.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f};
					for (int k = 0; k < 3; ++k) {
						if (tensor[index*30+k*10] > 0.5f) {
							dti[1] += tensor[index*30+k*10+4];
							dti[2] += tensor[index*30+k*10+5];
							dti[3] += tensor[index*30+k*10+6];
							dti[4] += tensor[index*30+k*10+7];
							dti[5] += tensor[index*30+k*10+8];
							dti[6] += tensor[index*30+k*10+9];
						}
					}
					pTempTensor->SetDTI(x, y, z, dti);
				}
			}
		}
		pTempTensor->SaveDTIFile(tempPathname);
		SafeDelete(pTempTensor);
	}
	
	SafeDeleteArray(tensor);
}
